TypeTree support in autodiff #144197

KMJ-007 · 2025-07-19T22:18:34Z

TypeTrees for Autodiff

What are TypeTrees?

Memory layout descriptors for Enzyme. Tell Enzyme exactly how types are structured in memory so it can compute derivatives efficiently.

Structure

TypeTree(Vec<Type>)

Type {
    offset: isize,  // byte offset (-1 = everywhere)
    size: usize,    // size in bytes
    kind: Kind,     // Float, Integer, Pointer, etc.
    child: TypeTree // nested structure
}

Example: `fn compute(x: &f32, data: &[f32]) -> f32`

Input 0: x: &f32

TypeTree(vec![Type {
    offset: 0, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: 0, size: 4, kind: Float,
        child: TypeTree::new()
    }])
}])

Input 1: data: &[f32]

TypeTree(vec![Type {
    offset: 0, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: -1, size: 4, kind: Float,  // -1 = all elements
        child: TypeTree::new()
    }])
}])

Output: f32

TypeTree(vec![Type {
    offset: 0, size: 4, kind: Float,
    child: TypeTree::new()
}])

Why Needed?

Enzyme can't deduce complex type layouts from LLVM IR
Prevents slow memory pattern analysis
Enables correct derivative computation for nested structures
Tells Enzyme which bytes are differentiable vs metadata

What Enzyme Does With This Information:

Without TypeTrees (current state):

; Enzyme sees generic LLVM IR:

define float @distance(i8* %p1, i8* %p2) {

; Has to guess what these pointers point to

; Slow analysis of all memory operations

; May miss optimization opportunities

}

With TypeTrees (our goal):

// Enzyme knows:

// - %p1 points to struct with f32 at +0, f32 at +4, i32 at +8

// - Only the f32 fields need derivatives

// - Can generate efficient derivative code directly

TypeTrees - Offset and -1 Explained

Type Structure

Type {

offset: isize, // WHERE this type starts

size: usize, // HOW BIG this type is

kind: Kind, // WHAT KIND of data (Float, Int, Pointer)

child: TypeTree // WHAT'S INSIDE (for pointers/containers)

}

Offset Values

Regular Offset (0, 4, 8, etc.)

Specific byte position within a structure

struct Point {

x: f32, // offset 0, size 4

y: f32, // offset 4, size 4

id: i32, // offset 8, size 4

}

TypeTree for &Point:

TypeTree(vec![

	Type { offset: 0, size: 4, kind: Float }, // x at byte 0

	Type { offset: 4, size: 4, kind: Float }, // y at byte 4

	Type { offset: 8, size: 4, kind: Integer } // id at byte 8

])

Offset -1 (Special: "Everywhere")

Means "this pattern repeats for ALL elements"

Example 1: Array `[f32; 100]`

TypeTree(vec![Type {

offset: -1, // ALL positions

size: 4, // each f32 is 4 bytes

kind: Float, // every element is float

}])

Instead of listing 100 separate Types with offsets 0,4,8,12...396

Example 2: Slice `&[i32]`

// Pointer to slice data

TypeTree(vec![Type {

	offset: 0, size: 8, kind: Pointer,

	child: TypeTree(vec![Type {

	offset: -1, // ALL slice elements

	size: 4, // each i32 is 4 bytes

	kind: Integer

	}])

}])

Example 3: Mixed Structure

struct Container {

	header: i64, // offset 0

	data: [f32; 1000], // offset 8, but elements use -1

}

TypeTree(vec![

	Type { offset: 0, size: 8, kind: Integer }, // header

	Type { offset: 8, size: 4000, kind: Pointer,

	child: TypeTree(vec![Type {

	offset: -1, size: 4, kind: Float // ALL array elements

}])

}

])

KMJ-007 · 2025-07-19T23:50:35Z

Currently, I have implemented only for memcpy

KMJ-007 · 2025-07-19T23:50:53Z

r? @ZuseZ4

rustbot · 2025-07-19T23:51:03Z

Some changes occurred in compiler/rustc_ast/src/expand/autodiff_attrs.rs

cc @ZuseZ4

Some changes occurred in compiler/rustc_codegen_llvm/src/builder/autodiff.rs

cc @ZuseZ4

Some changes occurred in compiler/rustc_codegen_ssa

cc @WaffleLapkin

Some changes occurred in compiler/rustc_monomorphize/src/partitioning/autodiff.rs

cc @ZuseZ4

rustbot · 2025-07-20T06:10:39Z

Some changes occurred in compiler/rustc_codegen_gcc

cc @antoyo, @GuillaumeGomez

KMJ-007 · 2025-07-21T12:01:35Z

CI is failing, fixing them!

compiler/rustc_codegen_gcc/src/builder.rs

rustbot · 2025-07-23T17:33:55Z

Some changes occurred in src/tools/enzyme

cc @ZuseZ4

bors · 2025-08-06T08:52:37Z

☔ The latest upstream changes (presumably #143684) made this pull request unmergeable. Please resolve the merge conflicts.

KMJ-007 · 2025-08-09T06:58:46Z

tests/run-make/autodiff/type-trees/type-analysis/memcpy/rmake.rs

+use run_make_support::{llvm_filecheck, rfs, rustc};
+
+fn main() {
+    // First, compile to LLVM IR to check for enzyme_type attributes
+    let _ir_output = rustc()
+        .input("memcpy.rs")
+        .arg("-Zautodiff=Enable")
+        .arg("-Zautodiff=NoPostopt")
+        .opt_level("3")
+        .arg("-Clto=fat")
+        .arg("--emit=llvm-ir")
+        .arg("-o")
+        .arg("main.ll")
+        .run();
+
+    // Then compile with TypeTree analysis output for the existing checks
+    let output = rustc()
+        .input("memcpy.rs")
+        .arg("-Zautodiff=Enable,PrintTAFn=test_memcpy")
+        .arg("-Zautodiff=NoPostopt")
+        .opt_level("3")
+        .arg("-Clto=fat")
+        .arg("-g")
+        .run();
+
+    let stdout = output.stdout_utf8();
+    let stderr = output.stderr_utf8();
+    let ir_content = rfs::read_to_string("main.ll");
+
+    rfs::write("memcpy.stdout", &stdout);
+    rfs::write("memcpy.stderr", &stderr);
+    rfs::write("main.ir", &ir_content);
+
+    llvm_filecheck().patterns("memcpy.check").stdin_buf(stdout).run();
+
+    llvm_filecheck().patterns("memcpy-ir.check").stdin_buf(ir_content).run();


@jieyouxu can you review this run make test, i am creating two check files to test IR and type analysis from enzyme, is this correct way or i should be combining them in one file with one single check?

Signed-off-by: Karan Janthe <[email protected]>

bors · 2025-08-12T17:08:15Z

☔ The latest upstream changes (presumably #145300) made this pull request unmergeable. Please resolve the merge conflicts.

ZuseZ4 · 2025-08-13T21:09:06Z

compiler/rustc_middle/src/ty/mod.rs

@@ -2194,3 +2201,335 @@ pub struct DestructuredConst<'tcx> {
    pub variant: Option<VariantIdx>,
    pub fields: &'tcx [ty::Const<'tcx>],
 }
+
+// Some types are used a lot. Make sure they don't unintentionally get bigger.
+#[cfg(target_pointer_width = "64")]


That's probably code from someone else, can we drop it?

ZuseZ4 · 2025-08-13T21:10:43Z

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

+}
+
+#[cfg(not(llvm_enzyme))]
+#[allow(dead_code)]


compiler/rustc_codegen_gcc/src/builder.rs

ZuseZ4 · 2025-08-13T21:15:29Z

compiler/rustc_resolve/src/lib.rs

@@ -1253,6 +1253,9 @@ pub struct Resolver<'ra, 'tcx> {
    // that were encountered during resolution. These names are used to generate item names
    // for APITs, so we don't want to leak details of resolution into these names.
    impl_trait_names: FxHashMap<NodeId, Symbol>,
+
+    /// Mapping of autodiff function IDs
+    autodiff_map: FxHashMap<LocalDefId, LocalDefId>,


Do we actually use the autodiff_map anywhere?

ZuseZ4 · 2025-08-13T21:32:09Z

compiler/rustc_middle/src/ty/mod.rs

+                _ => panic!(""),
+            };
+
+            let fields = adt_def.all_fields();


let fields = adt_def .all_fields() .into_iter() .zip(offsets.into_iter())

ZuseZ4 · 2025-08-13T21:35:09Z

compiler/rustc_middle/src/ty/mod.rs

+            if inner_ty.is_slice() {
+                // We know that the length will be passed as extra arg.
+                let child = typetree_from_ty(inner_ty, tcx, 1, safety, &mut visited, span);
+                let tt = Type { offset: -1, kind: Kind::Pointer, size: 8, child };


I think on 32 bit targets we probably have size: 4 here, can you check that (e.g. on the playground) and make it selective if needed?

ZuseZ4 · 2025-08-13T21:37:49Z

compiler/rustc_middle/src/ty/mod.rs

+    FncTree { args, ret }
+}
+
+fn typetree_from_ty<'a>(


Iirc I added the visited Vec here to detect cycles (recursive types). Can you verify that and add a comment describing what it does?

ZuseZ4 · 2025-08-13T21:38:13Z

compiler/rustc_middle/src/ty/mod.rs

+            assert!(span.is_some());
+            let span = span.unwrap();
+
+            tcx.sess.dcx().emit_warn(AutodiffUnsafeInnerConstRef { span, ty });


This should also have a test

ZuseZ4 · 2025-08-13T21:38:50Z

compiler/rustc_middle/src/ty/mod.rs

+            match ty {
+                x if x == tcx.types.f32 => (Kind::Float, 4),
+                x if x == tcx.types.f64 => (Kind::Double, 8),
+                _ => panic!("floatTy scalar that is neither f32 nor f64"),


I think we support f16 and f128 now, can you add those (And some tests)?

ZuseZ4 · 2025-08-13T21:52:35Z

compiler/rustc_middle/src/ty/mod.rs

+//    Not an error, because it only causes issues if they are actually read, which we don't check
+//    yet. We should add such analysis to relibably either issue an error or accept without warning.
+//    If there only were some research to do that...
+pub fn fnc_typetrees<'tcx>(


This fn is comparably complex, so we should have tests which cover every case handled here, or you split some of the logic out in follow-up PRs.

I would probably squash the current PR into 1-3 commits to start and only add a very minimal handling to fnc_typtrees (and other functions in this file). Then you can add follow-up commits for the extra handling. E.g. one commit for array handling + array tests, one commit for simd and the simd tests, one for recursive handling, and one for the InnerConstRef handling, etc. This way we can merge it incrementally and we're sure that every piece in this function actually works and is tested.

Yes, I was thinking about the same, should i do follow PRs adding new handling and tests or we have some basic typetree requirement which this PR must complete, so this can be merged.

What should be goal of the current PR ? basic typetree support, working memcpy with tests and i think NOTT flag, right?

ZuseZ4 · 2025-08-13T22:15:39Z

You should also add a Flag which disables all TypeTree additions (in case that it causes bugs, or just for A/B testing to see where the typetree's allow us to compile something, or have a compile time impact). We already have autodiff=Enable, just add another option to that enum, so people can pass -Zautodiff=Enable,NoTT.

rustbot added F-autodiff `#![feature(autodiff)]` S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. labels Jul 19, 2025

KMJ-007 mentioned this pull request Jul 19, 2025

[WIP] TypeTree support in autodiff #143490

Closed

This comment has been minimized.

Sign in to view

rustbot added the A-LLVM Area: Code generation parts specific to LLVM. Both correctness bugs and optimization-related issues. label Jul 19, 2025

This comment has been minimized.

Sign in to view

rustbot assigned ZuseZ4 Jul 19, 2025

KMJ-007 marked this pull request as ready for review July 19, 2025 23:50

rustbot added S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. and removed S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. labels Jul 19, 2025

This comment has been minimized.

Sign in to view

antoyo reviewed Jul 21, 2025

View reviewed changes

compiler/rustc_codegen_gcc/src/builder.rs Outdated Show resolved Hide resolved

This comment has been minimized.

Sign in to view

rustbot added has-merge-commits PR has merge commits, merge with caution. S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. labels Jul 23, 2025

rust-cloud-vms bot force-pushed the type-tree branch from d3f78d7 to 94aa8bc Compare July 23, 2025 17:37

This comment has been minimized.

Sign in to view

rustbot removed the has-merge-commits PR has merge commits, merge with caution. label Jul 23, 2025

rust-cloud-vms bot force-pushed the type-tree branch from 83aa06f to a28f311 Compare August 9, 2025 06:44

This comment has been minimized.

Sign in to view

KMJ-007 commented Aug 9, 2025

View reviewed changes

KMJ-007 requested a review from ZuseZ4 August 9, 2025 06:59

KMJ-007 added 20 commits August 9, 2025 13:22

sync with main

13e57dd

Signed-off-by: Karan Janthe <[email protected]>

memcpy compiling

1759277

Signed-off-by: Karan Janthe <[email protected]>

fmt

70d5ff8

Signed-off-by: Karan Janthe <[email protected]>

typoes fixing

73f2e2f

Signed-off-by: Karan Janthe <[email protected]>

typo

ccab2da

Signed-off-by: Karan Janthe <[email protected]>

fixing todo issue

6660dd6

Signed-off-by: Karan Janthe <[email protected]>

trailing whitespace

5f4cf00

Signed-off-by: Karan Janthe <[email protected]>

fixing white space issue because of docs comments

778f780

Signed-off-by: Karan Janthe <[email protected]>

fixing gcc memcpy

a7e256b

Signed-off-by: Karan Janthe <[email protected]>

wip, fixing ci

346bc6e

formatting

98765c5

Signed-off-by: Karan Janthe <[email protected]>

fixing ffi issue

019f248

noinline wait for sometimes

73b6dbb

fixing ir

00940cb

Signed-off-by: Karan Janthe <[email protected]>

fixing ci

c730df0

Signed-off-by: Karan Janthe <[email protected]>

fixing ci

44db42f

Signed-off-by: Karan Janthe <[email protected]>

pr review changes

28c4faa

fixing typetree metadata issue

7374020

Signed-off-by: Karan Janthe <[email protected]>

ci changes

db9c40c

Signed-off-by: Karan Janthe <[email protected]>

fixing ci changes

cffae78

Signed-off-by: Karan Janthe <[email protected]>

rust-cloud-vms bot force-pushed the type-tree branch from 110292b to cffae78 Compare August 9, 2025 13:23

ZuseZ4 reviewed Aug 13, 2025

View reviewed changes

TypeTree support in autodiff #144197

Are you sure you want to change the base?

TypeTree support in autodiff #144197

Conversation

KMJ-007 commented Jul 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

TypeTrees for Autodiff

What are TypeTrees?

Structure

Example: fn compute(x: &f32, data: &[f32]) -> f32

Why Needed?

What Enzyme Does With This Information:

TypeTrees - Offset and -1 Explained

Type Structure

Offset Values

Regular Offset (0, 4, 8, etc.)

Offset -1 (Special: "Everywhere")

Example 1: Array [f32; 100]

Example 2: Slice &[i32]

Example 3: Mixed Structure

Uh oh!

This comment has been minimized.

This comment has been minimized.

KMJ-007 commented Jul 19, 2025

Uh oh!

KMJ-007 commented Jul 19, 2025

Uh oh!

rustbot commented Jul 19, 2025

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

rustbot commented Jul 20, 2025

Uh oh!

This comment has been minimized.

KMJ-007 commented Jul 21, 2025

Uh oh!

Uh oh!

rustbot commented Jul 23, 2025

Uh oh!

This comment has been minimized.

This comment has been minimized.

bors commented Aug 6, 2025

Uh oh!

This comment has been minimized.

Choose a reason for hiding this comment

Uh oh!

bors commented Aug 12, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ZuseZ4 commented Aug 13, 2025

Uh oh!

Uh oh!

KMJ-007 commented Jul 19, 2025 •

edited

Loading

Example: `fn compute(x: &f32, data: &[f32]) -> f32`

Example 1: Array `[f32; 100]`

Example 2: Slice `&[i32]`